package org.example.myleet.rosalind.mend;

import org.example.myleet.rosalind.utils.NewickTreeNode;
import org.example.myleet.rosalind.utils.NewickTreeUtil;

import java.io.*;
import java.math.BigDecimal;
import java.math.RoundingMode;

public class Solution {

    private static final BigDecimal HALF = BigDecimal.valueOf(0.5);

    public static void main(String[] args) throws IOException {
        Solution solution = new Solution();
        String path = solution.getClass().getResource("").getPath() + "example.txt";
        BufferedReader br = new BufferedReader(new InputStreamReader(new FileInputStream(path)));
        String source = br.readLine().trim();
        NewickTreeNode root = NewickTreeUtil.parse(null, source);
        BigDecimal[] result = dfsCalculateAlleleProbability(root);
        System.out.printf("%s %s %s\n", result[0].setScale(6, RoundingMode.HALF_UP), result[1].setScale(6, RoundingMode.HALF_UP), result[2].setScale(6, RoundingMode.HALF_UP));
    }
    private static BigDecimal[] dfsCalculateAlleleProbability(NewickTreeNode node) {
        //probabilities of organisms with AA/Aa/aa genotype
        BigDecimal[] result = new BigDecimal[3];
        if (null == node.getChildren() || node.getChildren().size() < 1) {
            //probability for each genotype
            if ("AA".equals(node.getValue())) {
                result[0] = BigDecimal.ONE;
                result[1] = BigDecimal.ZERO;
                result[2] = BigDecimal.ZERO;
            } else if ("Aa".equals(node.getValue()) || "aA".equals(node.getValue())) {
                result[0] = BigDecimal.ZERO;
                result[1] = BigDecimal.ONE;
                result[2] = BigDecimal.ZERO;
            } else {//aa
                result[0] = BigDecimal.ZERO;
                result[1] = BigDecimal.ZERO;
                result[2] = BigDecimal.ONE;
            }
            return result;
        }
        //Use Depth First Search to calculate the child's genotype probability
        BigDecimal[] first = dfsCalculateAlleleProbability(node.getChildren().get(0));
        BigDecimal[] second = dfsCalculateAlleleProbability(node.getChildren().get(1));
        //Calculate the dominant allele's probability and recessive allele's probability of the first parent
        BigDecimal firstDominant = first[0].add(first[1].multiply(HALF));
        BigDecimal firstRecessive = first[2].add(first[1].multiply(HALF));
        //Calculate the dominant allele's probability and recessive allele's probability of the second parent
        BigDecimal secondDominant = second[0].add(second[1].multiply(HALF));
        BigDecimal secondRecessive = second[2].add(second[1].multiply(HALF));
        //Then their child's genotype(AA/Aa/aa) can be calculated from the dominant and recessive allele combination
        result[0] = firstDominant.multiply(secondDominant);
        result[1] = firstDominant.multiply(secondRecessive).add(firstRecessive.multiply(secondDominant));
        result[2] = firstRecessive.multiply(secondRecessive);
        return result;
    }
}
